import image_detect
import image_classification
import os
from Vehicle_Detecter_Win import Ui_Vehicle_Detecter
from PyQt5 import QtWidgets, QtGui, QtCore
from PyQt5.QtCore import QThread,Qt,pyqtSignal
from PyQt5.QtWidgets import QFileDialog
from PyQt5.QtGui import QIcon
from PyQt5.QtGui import QFont
from PIL import Image
import time
import platform
from enum import Enum
# 定义Vehicle_Detecter类，继承Ui_Vehicle_Detecter
class Vehicle_Detecter(QtWidgets.QWidget,Ui_Vehicle_Detecter):
    def __init__(self):
        super(Vehicle_Detecter,self).__init__()
        self.setupUi(self)
        #定义应用程序icon图标
        self.setWindowIcon(QIcon('res/img/icon.jpg'))
        self.output.setFont(QFont('simfang', 18))
        #定义按键操作接口
        self.SelectButton.clicked.connect(self.select_image)
        self.SelOutButton.clicked.connect(self.set_output_path)
        self.RunButton.clicked.connect(self.run)
        self.OpenDirButton.clicked.connect(self.open)
        #由于检测和分类耗时较长，开子线程操作
        self.detect_thread = detect_thread()
        #加载权重时间较长，开子线程操作
        self.load_thread = load_thread()
        #定义子线程回调
        self.detect_thread.sinOut.connect(self.display_result)
        self.load_thread.sinOut.connect(self.loaded)
        #定义开始检测时间
        self.start_time = 0

    # 初始化image_detect和image_classification,加载网络权重
    def load_pb(self):
        #显示加载图片和提示
        self.gif = QtGui.QMovie('./res/img/pre.gif')
        self.label.setMovie(self.gif)
        self.gif.start()
        p = QtGui.QPixmap()
        self.graphicsView.scene = QtWidgets.QGraphicsScene()  # 创建一个图片元素的对象
        item = QtWidgets.QGraphicsPixmapItem(p)  # 创建一个变量用于承载加载后的图片
        self.graphicsView.scene.addItem(item)  # 将加载后的图片传递给scene对象
        self.graphicsView.setScene(self.graphicsView.scene)
        self.output.setText('Loading......')
        self.load_thread.start()
    #加载完的回调
    def loaded(self):
        self.gif.stop()
        self.gif = QtGui.QMovie('')
        self.label.setMovie(self.gif)
        self.output.setText('')
    #选取图片按键接口，支持png,jpg,bmp三种格式
    def select_image(self):
        fileName, filetype = QFileDialog.getOpenFileName(self, "选取文件", "./","image Files (*.jpg *.png *.bmp)")
        #print(fileName)
        self.ImgPath.setText(fileName)
    #设置输出目录接口
    def set_output_path(self):
        directory = QFileDialog.getExistingDirectory(self, "选取文件夹", "./")  # 起始路径
        self.OutPath.setText(directory)
    #设置打开输出目录接口
    def open(self):
        out_path = self.OutPath.text().strip()
        #如果没有选择或输入，则打开默认输出目录
        if len(out_path) == 0:
            out_path = 'out'
            if not os.path.exists(out_path):
                os.makedirs(out_path)
        #如果输入的目录不存在，则在ui上提示
        if not os.path.exists(out_path):
            self.show_error(Except_Type.nodir, out_path)
        else:
            sysstr = platform.system()
            if(sysstr == "Windows"):
                os.startfile(out_path)
            elif(sysstr == "Linux"):
                cmd = 'nautilus ' + out_path
                os.system(cmd)
            else:
                print("Not supported system os".format(sysstr))
    #检测接口
    def run(self):
        #检测图片是否存在，不存在则提示并退出
        file_path = self.ImgPath.text()
        if not os.path.exists(file_path):
            self.show_error(Except_Type.nofile, file_path)
            return
        dirname, filename = os.path.split(file_path)
        fname, fename = os.path.splitext(filename)
        if fename not in ['.jpg', '.bmp', '.png']:
            self.show_error(Except_Type.wrongtype, filename)
            return
        try:
            image = Image.open(file_path)
        except IOError as e:
            self.show_error(Except_Type.brokenfile, filename)
            return
        image.close()
        out_path = self.OutPath.text().strip()
        #创建输出目录，如果无法创建，则提示并退出
        if len(out_path) == 0:
            out_path = 'out'
        if not os.path.exists(out_path):
            try:
                os.makedirs(out_path)
            except FileNotFoundError:
                self.show_error(Except_Type.createdirfail, out_path)
                return
        #if not os.path.exists(out_path):
        #    self.show_error(Except_Type.createdirfail, out_path)
        #    return
        #记录开始时间
        self.start_time = time.time()
        #播放等待图片
        self.gif = QtGui.QMovie('res/img/pre.gif')
        self.label.setMovie(self.gif)
        self.gif.start()
        p = QtGui.QPixmap()
        self.graphicsView.scene = QtWidgets.QGraphicsScene()
        item = QtWidgets.QGraphicsPixmapItem(p)
        self.graphicsView.scene.addItem(item)
        self.graphicsView.setScene(self.graphicsView.scene)
        self.output.setText('Please waiting......')
        #选择检测方式
        if self.FastButton.isChecked():
            model = 'ssd'
        elif self.NormalButton.isChecked():
            model = 'frnn'
        elif self.EndButton.isChecked():
            model = 'direct'
        #运行子进程进行检测
        self.detect_thread.set_path(file_path, out_path, model)
        self.detect_thread.start()
        #dic = image_detect.detecter(file_path, model, out_path)
        #print(dic)
    '''显示检测结果
        result:检测得到的结果组成的字典'''
    def display_result(self, result):
        #停止等待图片播放
        self.gif.stop()
        self.gif = QtGui.QMovie('')
        self.label.setMovie(self.gif)
        #设置输出文字，包括文件名、检测耗时、检测到的车辆型号及概率
        file_path = self.ImgPath.text()
        dirname, filename = os.path.split(file_path)
        pass_time = time.time() - self.start_time
        text = 'Name:  ' +filename + '\n\n' + 'Cost: ' + "%.3f s" % pass_time + '\n\n'
        #如果没检测到车，输出没有车的结果
        if len(result[file_path]) == 0:
            self.show_error(Except_Type.novehicle, text)
            p = QtGui.QPixmap(file_path).scaled(940, 590, QtCore.Qt.KeepAspectRatio)
        else:
            #检测到的车辆按概率从大到小排列显示
            for i in range(len(result[file_path])//3):

                text = text + 'Model: ' +  \
                       str(image_detect.get_class_name(result[file_path][i * 3])) + '\n' \
                       + 'Score: ' + "%.2f%%" % (result[file_path][i * 3 + 1] * 100) + '\n'
            self.output.setText(text)
            #保存输出图片到out目录
            outimg = image_detect.save_result(result)
            #对显示的图片进行缩放
            with Image.open(outimg) as img:
                if img.size[0] > 2000 or img.size[1] > 1100:
                    p = QtGui.QPixmap(outimg)
                else:
                    p = QtGui.QPixmap(outimg).scaled(940, 590, QtCore.Qt.KeepAspectRatio)
        #在程序中显示结果
        self.graphicsView.scene = QtWidgets.QGraphicsScene()
        item = QtWidgets.QGraphicsPixmapItem(p)
        self.graphicsView.scene.addItem(item)
        self.graphicsView.setScene(self.graphicsView.scene)
        #删除产生的临时文件，windows和ubuntu兼容
        sysstr = platform.system()
        if(sysstr == "Windows"):
            #lists = os.listdir('.tmp')
            #for file in lists:
            #    os.remove(os.path.join('.tmp', file))
            pass
        elif(sysstr == "Linux"):
            os.system('rm -rf .tmp/*')
        else:
            print("Not supported system os".format(sysstr))
    '''异常提示处理函数
    except_code:异常编号
    msg：异常附加信息'''
    def show_error(self, except_code, msg):
        if except_code == Except_Type.nofile:
            self.output.setText("File {} is not exist!".format(msg))
        elif except_code == Except_Type.nodir:
            self.output.setText("Dir {} is not exist!".format(msg))
        elif except_code == Except_Type.createdirfail:
            self.output.setText("Dir {} can not be created!".format(msg))
        elif except_code == Except_Type.wrongtype:
            self.output.setText("File {} type do not support!".format(msg))
        elif except_code == Except_Type.brokenfile:
            self.output.setText("File {} has been damaged!".format(msg))
        elif except_code == Except_Type.novehicle:
            self.output.setText("{} No vehicle!".format(msg))
        elif except_code == Except_Type.other:
            self.output.setText("Unknow Except!")
#定义异常类型
class Except_Type(Enum):
    nofile = 1
    nodir = 2
    wrongtype = 3
    brokenfile = 4
    createdirfail = 5
    novehicle = 6
    other = 7


'''用于检测的子进程类'''
class detect_thread(QThread):
    #定义回调参数为dict
    sinOut = pyqtSignal(dict)
    def __init__(self):
        super().__init__()
    '''主线程到子线程参数传递方法
        file_path:输入图片全路径
        out_path：输出目录
        model：检测模型选择'''
    def set_path(self,file_path, out_path, model):
        self.file_path = file_path
        self.out_path = out_path
        self.model = model
    #运行检测的函数入口
    def run(self):
        result = image_detect.detecter(self.file_path, self.model, self.out_path)
        #检测完后返回结果到主线程
        self.sinOut.emit(result)
#load 子线程
class load_thread(QThread):
    sinOut = pyqtSignal()
    def __init__(self):
        super().__init__()
    #子线程主体
    def run(self):
        image_detect.init()
        image_classification.init()
        self.sinOut.emit()
if __name__ == '__main__':
    import sys
    app = QtWidgets.QApplication(sys.argv)
    MainWindow = Vehicle_Detecter()
    MainWindow.show()
    MainWindow.load_pb()
    sys.exit(app.exec_())
